import os
import gradio as gr
from model import GLMModel, OpenAIModel
from translator import PDFTranslator

work_dir = os.getcwd()

def translator(model_type, model_name, model_url, openai_key, file_path, file_format):
    model = None
    if model_type == 'OpenAIModel':
        model = OpenAIModel(model=model_name, api_key=openai_key)
    elif model_type == 'GLMModel':
        model = GLMModel(model_url, 300)

    task = PDFTranslator(model)
    out_file_path = ''
    if file_format == 'pdf':
        out_file_path = os.path.join(work_dir, 'out.pdf')
    elif file_format == 'markdown':
        out_file_path = os.path.join(work_dir, 'out.md')

    task.translate_pdf(file_path, file_format, output_file_path=out_file_path)
    return out_file_path

webui = gr.Interface(fn=translator,
                     inputs=[
                         gr.Radio(['OpenAIModel', 'GLMModel']),
                         'text', 'text', 'text', 'file',
                         gr.Radio(['pdf', 'markdown'])
                     ],
                     outputs=['text']
                     )
webui.launch()